import spconv.pytorch as spconv
import numpy as np
import torch.nn as nn

class ResSPConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, indice_key = "key", kernel_size = 3):
        super().__init__()
        self.spconv1 = spconv.SubMConv3d(in_channels=in_channels, out_channels=out_channels, \
                                               stride = 1, kernel_size=kernel_size, indice_key = indice_key)
        self.spconv2 = spconv.SubMConv3d(in_channels=out_channels, out_channels=out_channels, \
                                               stride = 1, kernel_size=kernel_size, indice_key = indice_key)
        self.instance_norm1 = nn.LayerNorm(out_channels)
        self.instance_norm2 = nn.LayerNorm(out_channels)
        
        self.relu = nn.ReLU()
        self.resnetblock = spconv.SparseSequential(
            self.spconv1,
            self.instance_norm1,
            self.relu,
            self.spconv2,
            self.instance_norm2
        )

    def forward(self, voxel_dict, spatial_shape, batch_size):
        features = voxel_dict['features']
        coordinates = voxel_dict['coordinates']
        input_sparse = spconv.SparseConvTensor(features, coordinates, spatial_shape, batch_size)
        residual = input_sparse
        out = self.resnetblock(input_sparse)
        out = out + residual
        output = self.relu(out.features)
        return output, out.indices

class ConvFFN(nn.Module):
    def __init__(self, d_model, dim_feedforward, mlp_dropout=0, kernel_size = 3, indice_key = "key"):
        super().__init__()
        self.spconv1 = spconv.SubMConv3d(in_channels=dim_feedforward, out_channels=dim_feedforward, \
                                               kernel_size=kernel_size, indice_key = indice_key)
        # Implementation of Feedforward model
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(mlp_dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)
        self.d_model = d_model
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Identity()
        self.gelu = nn.functional.gelu
        self.relu = nn.functional.relu

    def forward(self, src, voxel_coords, spatial_shape, batch_size):
        src2 = self.relu(self.linear1(src)).float()
        input_sparse = spconv.SparseConvTensor(src2, voxel_coords, spatial_shape, batch_size)
        src2 = self.spconv1(input_sparse).features
        src2 = self.linear2(self.dropout(self.gelu(src2)))
        src = src + self.dropout2(src2)
        src = self.norm2(src)
        return src
    
